{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### TAO remote client (object detection with YOLO)\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "![image](https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png)\n",
    "\n",
    "\n",
    "### The workflow in a nutshell\n",
    "\n",
    "- Model Actions\n",
    "    - Prune, retrain\n",
    "    - Export\n",
    "    - Convert\n",
    "    - Inference on TRT\n",
    "\n",
    "### Table of contents\n",
    "\n",
    "1. [Provide FP32 export specs](#head-1)\n",
    "1. [Run FP32 export](#head-2)\n",
    "1. [Provide model convert specs](#head-3)\n",
    "1. [Run model convert](#head-4)\n",
    "1. [Provide TAO inference specs](#head-5)\n",
    "1. [Run TAO inference](#head-6)\n",
    "1. [Provide prune specs](#head-7)\n",
    "1. [Run prune](#head-8)\n",
    "1. [Provide retrain specs](#head-9)\n",
    "1. [Run retrain](#head-10)\n",
    "1. [Provide evaluate specs](#head-11)\n",
    "1. [Run evaluate on retrain](#head-12)\n",
    "1. [Provide Int8 export specs](#head-13)\n",
    "1. [Run Int8 export](#head-14)\n",
    "1. [Provide model convert specs](#head-15)\n",
    "1. [Run model convert](#head-16)\n",
    "1. [Provide TAO inference specs](#head-17)\n",
    "1. [Run TAO inference](#head-18)\n",
    "1. [Delete experiment](#head-19)\n",
    "1. [Delete datasets](#head-20)\n",
    "1. [Unmount shared volume](#head-21)\n",
    "\n",
    "### Requirements\n",
    "Please find the server requirements [here](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_setup.html#)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import subprocess\n",
    "import getpass\n",
    "import uuid\n",
    "import json\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "namespace = 'default'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### FIXME\n",
    "\n",
    "1. Set model_id in FIXME 1\n",
    "1. Set train_job_id in FIXME 2\n",
    "1. Assign the ip_address and port_number in FIXME 3 and FIXME 4 ([info](https://docs.nvidia.com/tao/tao-toolkit/text/tao_toolkit_api/api_rest_api.html))\n",
    "1. Set NGC API key in FIXME 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = FIXME #FIXME1\n",
    "train_job_id = FIXME #FIXME2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Set the remote service base URL <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Define the node_addr and port number\n",
    "node_addr = \"<ip_address>\" # FIXME3 example: 10.137.149.22\n",
    "node_port = \"<port_number>\" # FIXME4 example: 32334\n",
    "# In host machine, node ip_address and port number can be obtained as follows,\n",
    "# ip_address: hostname -i\n",
    "# port_number: kubectl get service ingress-nginx-controller -o jsonpath='{.spec.ports[0].nodePort}'\n",
    "%env BASE_URL=http://{node_addr}:{node_port}/{namespace}/api/v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# FIXME: Set ngc_api_key valiable\n",
    "ngc_api_key = \"<ngc_api_key>\" # FIXME5 example: zZYtczM5amdtdDcwNjk0cnA2bGU2bXQ3bnQ6NmQ4NjNhMDItMTdmZS00Y2QxLWI2ZjktNmE5M2YxZTc0OGyM\n",
    "\n",
    "# Exchange NGC_API_KEY for JWT\n",
    "identity = json.loads(subprocess.getoutput(f'nvtl login --ngc-api-key {ngc_api_key}'))\n",
    "\n",
    "%env USER={identity['user_id']}\n",
    "%env TOKEN={identity['token']}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "user = getpass.getuser()\n",
    "home = os.path.expanduser('~')\n",
    "\n",
    "! echo \"Password for {user}\"\n",
    "password = getpass.getpass()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def my_tail(logs_dir, log_file):\n",
    "    %env LOG_FILE={logs_dir}/{log_file}\n",
    "    ! mkdir -p {logs_dir}\n",
    "    ! [ ! -f \"$LOG_FILE\" ] && touch $LOG_FILE && chmod 666 $LOG_FILE\n",
    "    ! tail -f -n +1 $LOG_FILE | while read LINE; do echo \"$LINE\"; [[ \"$LINE\" == \"EOF\" ]] && pkill -P $$ tail; done"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "logs_dir = os.path.join(home, 'shared', 'users', os.environ['USER'], 'models', model_id, 'logs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Provide FP32 export specs <a class=\"anchor\" id=\"head-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default export model specs\n",
    "! nvtl yolo-v4 model-export-defaults --id {model_id} | tee ~/shared/users/{os.environ['USER']}/models/{model_id}/specs/export.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize export model specs\n",
    "specs_path = os.path.join(home, 'shared', 'users', os.environ['USER'], 'models', model_id, 'specs', 'export.json')\n",
    "\n",
    "with open(specs_path , \"r\") as specs_file:\n",
    "    specs = json.load(specs_file)\n",
    "\n",
    "specs[\"data_type\"] = \"fp32\"\n",
    "\n",
    "with open(specs_path, \"w\") as specs_file:\n",
    "    json.dump(specs, specs_file, indent=2)\n",
    "\n",
    "print(json.dumps(specs, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run FP32 export <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "fp32_export_job_id = subprocess.getoutput(f\"nvtl yolo-v4 model-export --id {model_id} --job {train_job_id}\")\n",
    "print(fp32_export_job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "log_file = f\"{fp32_export_job_id}.txt\"\n",
    "my_tail(logs_dir, log_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Provide model convert specs <a class=\"anchor\" id=\"head-3\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default convert model specs\n",
    "! nvtl yolo-v4 model-convert-defaults --id {model_id} | tee ~/shared/users/{os.environ['USER']}/models/{model_id}/specs/convert.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize convert model specs\n",
    "specs_path = os.path.join(home, 'shared', 'users', os.environ['USER'], 'models', model_id, 'specs', 'convert.json')\n",
    "\n",
    "with open(specs_path , \"r\") as specs_file:\n",
    "    specs = json.load(specs_file)\n",
    "\n",
    "specs[\"t\"] = \"fp32\"\n",
    "specs[\"b\"] = 8\n",
    "specs[\"p\"] = \"Input,1x3x736x1280,8x3x736x1280,16x3x736x1280\"\n",
    "\n",
    "with open(specs_path, \"w\") as specs_file:\n",
    "    json.dump(specs, specs_file, indent=2)\n",
    "\n",
    "print(json.dumps(specs, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run model convert <a class=\"anchor\" id=\"head-4\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "convert_job_id = subprocess.getoutput(f\"nvtl yolo-v4 model-convert --id {model_id} --job {fp32_export_job_id}\")\n",
    "print(convert_job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "log_file = f\"{convert_job_id}.txt\"\n",
    "my_tail(logs_dir, log_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Provide TAO inference specs <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default inference model specs\n",
    "! nvtl yolo-v4 model-inference-defaults --id {model_id} | tee ~/shared/users/{os.environ['USER']}/models/{model_id}/specs/inference.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize TAO inference specs\n",
    "specs_path = os.path.join(home, 'shared', 'users', os.environ['USER'], 'models', model_id, 'specs', 'inference.json')\n",
    "\n",
    "with open(specs_path , \"r\") as specs_file:\n",
    "    specs = json.load(specs_file)\n",
    "\n",
    "specs[\"augmentation_config\"][\"output_width\"] = 1280 # Setting to the dataset's original resolution's width\n",
    "specs[\"augmentation_config\"][\"output_height\"] = 736 # Setting to the dataset's original resolution's height\n",
    "\n",
    "with open(specs_path, \"w\") as specs_file:\n",
    "    json.dump(specs, specs_file, indent=2)\n",
    "\n",
    "print(json.dumps(specs, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run TAO inference <a class=\"anchor\" id=\"head-6\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tao_inference_job_id = subprocess.getoutput(f\"nvtl yolo-v4 model-inference --id {model_id} --job {convert_job_id}\")\n",
    "print(tao_inference_job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "log_file = f\"{tao_inference_job_id}.txt\"\n",
    "start_time = time.time()\n",
    "my_tail(logs_dir, log_file)\n",
    "inference_time = time.time() - start_time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Time in seconds for inference on unoptimized model is \", inference_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "job_dir = f\"{home}/shared/users/{os.environ['USER']}/models/{model_id}/{tao_inference_job_id}\"\n",
    "from IPython.display import Image\n",
    "import glob\n",
    "sample_image = glob.glob(f\"{job_dir}/images_annotated/*.jpg\")[6]\n",
    "Image(filename=sample_image) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### Provide prune specs <a class=\"anchor\" id=\"head-7\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default prune model specs\n",
    "! nvtl yolo-v4 model-prune-defaults --id {model_id} | tee ~/shared/users/{os.environ['USER']}/models/{model_id}/specs/prune.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### Run prune <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "prune_job_id = subprocess.getoutput(f\"nvtl yolo-v4 model-prune --id {model_id} --job {train_job_id}\")\n",
    "print(prune_job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "log_file = f\"{prune_job_id}.txt\"\n",
    "my_tail(logs_dir, log_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prune_job_dir = f\"{home}/shared/users/{os.environ['USER']}/models/{model_id}/{prune_job_id}\"\n",
    "pruned_model_size = json.loads(subprocess.getoutput(f'stat -c \"%s\" {prune_job_dir}/model.tlt'))\n",
    "\n",
    "train_job_dir = f\"{home}/shared/users/{os.environ['USER']}/models/{model_id}/{train_job_id}\"\n",
    "original_train_model_size = json.loads(subprocess.getoutput(f'stat -c \"%s\" {train_job_dir}/weights/yolov4_resnet18_epoch_010.tlt'))\n",
    "\n",
    "print(f\"The original trained model size is {original_train_model_size} KB\")\n",
    "print(f\"The pruned model size is {pruned_model_size} KB\")\n",
    "print(f\"The pruned model is {round(original_train_model_size/pruned_model_size,1)}x smaller than the original model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### Provide retrain specs <a class=\"anchor\" id=\"head-9\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default retrain model specs\n",
    "! nvtl yolo-v4 model-retrain-defaults --id {model_id} | tee ~/shared/users/{os.environ['USER']}/models/{model_id}/specs/retrain.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize retrain model specs\n",
    "specs_path = os.path.join(home, 'shared', 'users', os.environ['USER'], 'models', model_id, 'specs', 'retrain.json')\n",
    "\n",
    "with open(specs_path , \"r\") as specs_file:\n",
    "    specs = json.load(specs_file)\n",
    "\n",
    "specs[\"training_config\"][\"num_epochs\"] = 10\n",
    "specs[\"dataset_config\"][\"image_extension\"] = \"jpg\" # Setting to the dataset's image_file extension type\n",
    "\n",
    "specs[\"augmentation_config\"][\"output_width\"] = 1280 # Setting to the dataset's original resolution's width\n",
    "specs[\"augmentation_config\"][\"output_height\"] = 736 # Setting to the dataset's original resolution's height\n",
    "\n",
    "with open(specs_path, \"w\") as specs_file:\n",
    "    json.dump(specs, specs_file, indent=2)\n",
    "\n",
    "print(json.dumps(specs, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### Run retrain <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "retrain_job_id = subprocess.getoutput(f\"nvtl yolo-v4 model-retrain --id {model_id} --job {prune_job_id}\")\n",
    "print(retrain_job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "log_file = f\"{retrain_job_id}.txt\"\n",
    "my_tail(logs_dir, log_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Provide evaluate specs <a class=\"anchor\" id=\"head-11\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default evaluate model specs\n",
    "! nvtl yolo-v4 model-evaluate-defaults --id {model_id} | tee ~/shared/users/{os.environ['USER']}/models/{model_id}/specs/evaluate.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize evaluate model specs\n",
    "specs_path = os.path.join(home, 'shared', 'users', os.environ['USER'], 'models', model_id, 'specs', 'evaluate.json')\n",
    "\n",
    "with open(specs_path , \"r\") as specs_file:\n",
    "    specs = json.load(specs_file)\n",
    "\n",
    "specs[\"dataset_config\"][\"image_extension\"] = \"jpg\" # Setting to the dataset's image_file extension type\n",
    "\n",
    "specs[\"augmentation_config\"][\"output_width\"] = 1280 # Setting to the dataset's original resolution's width\n",
    "specs[\"augmentation_config\"][\"output_height\"] = 736 # Setting to the dataset's original resolution's height\n",
    "\n",
    "with open(specs_path, \"w\") as specs_file:\n",
    "    json.dump(specs, specs_file, indent=2)\n",
    "\n",
    "print(json.dumps(specs, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run evaluate on retrained model <a class=\"anchor\" id=\"head-12\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "eval2_job_id = subprocess.getoutput(f\"nvtl yolo-v4 model-evaluate --id {model_id} --job {retrain_job_id}\")\n",
    "print(eval2_job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "log_file = f\"{eval2_job_id}.txt\"\n",
    "my_tail(logs_dir, log_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Provide Int8 export specs <a class=\"anchor\" id=\"head-13\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default export model specs\n",
    "! nvtl yolo-v4 model-export-defaults --id {model_id} | tee ~/shared/users/{os.environ['USER']}/models/{model_id}/specs/export.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize export model specs\n",
    "specs_path = os.path.join(home, 'shared', 'users', os.environ['USER'], 'models', model_id, 'specs', 'export.json')\n",
    "\n",
    "with open(specs_path , \"r\") as specs_file:\n",
    "    specs = json.load(specs_file)\n",
    "\n",
    "specs[\"data_type\"] = \"int8\"\n",
    "specs[\"batches\"] = 10\n",
    "specs[\"batch_size\"] = 16\n",
    "\n",
    "with open(specs_path, \"w\") as specs_file:\n",
    "    json.dump(specs, specs_file, indent=2)\n",
    "\n",
    "print(json.dumps(specs, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run Int8 export <a class=\"anchor\" id=\"head-14\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "int8_export_job_id = subprocess.getoutput(f\"nvtl yolo-v4 model-export --id {model_id} --job {retrain_job_id}\")\n",
    "print(int8_export_job_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "log_file = f\"{int8_export_job_id}.txt\"\n",
    "my_tail(logs_dir, log_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Provide model convert specs <a class=\"anchor\" id=\"head-15\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default convert model specs\n",
    "! nvtl yolo-v4 model-convert-defaults --id {model_id} | tee ~/shared/users/{os.environ['USER']}/models/{model_id}/specs/convert.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize convert model specs\n",
    "specs_path = os.path.join(home, 'shared', 'users', os.environ['USER'], 'models', model_id, 'specs', 'convert.json')\n",
    "\n",
    "with open(specs_path , \"r\") as specs_file:\n",
    "    specs = json.load(specs_file)\n",
    "\n",
    "specs[\"t\"] = \"int8\"\n",
    "specs[\"b\"] = 8\n",
    "specs[\"p\"] = \"Input,1x3x736x1280,8x3x736x1280,16x3x736x1280\"\n",
    "\n",
    "with open(specs_path, \"w\") as specs_file:\n",
    "    json.dump(specs, specs_file, indent=2)\n",
    "\n",
    "print(json.dumps(specs, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run model convert <a class=\"anchor\" id=\"head-16\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "convert_job_id_2 = subprocess.getoutput(f\"nvtl yolo-v4 model-convert --id {model_id} --job {int8_export_job_id}\")\n",
    "print(convert_job_id_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "log_file = f\"{convert_job_id_2}.txt\"\n",
    "my_tail(logs_dir, log_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Provide TAO inference specs <a class=\"anchor\" id=\"head-17\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Default inference model specs\n",
    "! nvtl yolo-v4 model-inference-defaults --id {model_id} | tee ~/shared/users/{os.environ['USER']}/models/{model_id}/specs/inference.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Customize TAO inference specs\n",
    "specs_path = os.path.join(home, 'shared', 'users', os.environ['USER'], 'models', model_id, 'specs', 'inference.json')\n",
    "\n",
    "with open(specs_path , \"r\") as specs_file:\n",
    "    specs = json.load(specs_file)\n",
    "\n",
    "specs[\"augmentation_config\"][\"output_width\"] = 1280 # Setting to the dataset's original resolution's width\n",
    "specs[\"augmentation_config\"][\"output_height\"] = 736 # Setting to the dataset's original resolution's height\n",
    "\n",
    "with open(specs_path, \"w\") as specs_file:\n",
    "    json.dump(specs, specs_file, indent=2)\n",
    "\n",
    "print(json.dumps(specs, indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run TAO inference <a class=\"anchor\" id=\"head-18\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tao_inference_job_id_2 = subprocess.getoutput(f\"nvtl yolo-v4 model-inference --id {model_id} --job {convert_job_id_2}\")\n",
    "print(tao_inference_job_id_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Check status (the file won't exist until the backend Toolkit container is running -- can take several minutes)\n",
    "log_file = f\"{tao_inference_job_id_2}.txt\"\n",
    "start_time = time.time()\n",
    "my_tail(logs_dir, log_file)\n",
    "inference_time = time.time() - start_time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Time in seconds for inference on optimized model is \", inference_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "job_dir = f\"{home}/shared/users/{os.environ['USER']}/models/{model_id}/{tao_inference_job_id_2}\"\n",
    "from IPython.display import Image\n",
    "import glob\n",
    "sample_image = glob.glob(f\"{job_dir}/images_annotated/*.jpg\")[6]\n",
    "Image(filename=sample_image) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### Delete experiment <a class=\"anchor\" id=\"head-19\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! rm -rf ~/shared/users/{os.environ['USER']}/models/{model_id}\n",
    "! echo DONE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "MD"
    }
   },
   "source": [
    "### Delete datasets <a class=\"anchor\" id=\"head-20\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "datalore": {
     "hide_input_from_viewers": false,
     "hide_output_from_viewers": false,
     "type": "CODE"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! rm -rf ~/shared/users/{os.environ['USER']}/datasets/{train_dataset_id}\n",
    "! rm -rf ~/shared/users/{os.environ['USER']}/datasets/{eval_dataset_id}\n",
    "! rm -rf ~/shared/users/{os.environ['USER']}/datasets/{infer_dataset_id}\n",
    "! echo DONE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Unmount shared volume <a class=\"anchor\" id=\"head-21\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "command = \"umount ~/shared\"\n",
    "! echo {password} | sudo -S -k {command} && echo DONE"
   ]
  }
 ],
 "metadata": {
  "datalore": {
   "base_environment": "default",
   "computation_mode": "JUPYTER",
   "package_manager": "pip",
   "packages": [],
   "version": 1
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
